#include "typed_ast.hpp"
#include "symtable.hpp"
#include "ast.hpp"
#include "type.hpp"
#include "lexer.hpp"
#include "parser.hpp"

#include <cassert>
#include <iostream>
using namespace std;

TypedVisitor::TypedVisitor(SymbolTable* symtable) : _sym_table(symtable) {
    _has_error = false;
    _is_func = false;
}

TypedVisitor::~TypedVisitor() {
}

TypeTable* TypedVisitor::type_table() {
    return _sym_table->type_table();
}

Node* TypedVisitor::visit(Node* n) {
    return n->accept(this);
}

Node* TypedVisitor::visit(PrintNode* n) {
    return new PrintNode(visit(n->value()));
}

Node* TypedVisitor::visit(ReturnNode* n) {
    visit(n->value());
    if (!_sym_table->is_func()) {
        cout << "Can not use return outside of function" << endl;
        _has_error = true;
        return n;
    }
    n->set_type(n->value()->get_type(type_table()));
    return n;
}

// reuse old node.
Node* TypedVisitor::visit(BinaryOp* op) {
    Node* left = visit(op->left());
    Node* right = visit(op->right());
    op->set_left(left);
    op->set_right(right);
    op->set_type(op->left()->get_type(type_table()));

    return op;
}

Node* TypedVisitor::visit(ListNode* n) {
    ListNode * list_node = new ListNode();

    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        Node * new_node = visit(*it);
        if (new_node) {
            list_node->add(new_node);
        }
    }
    return list_node;
}

Node* TypedVisitor::visit(StmtsNode* n) {
    StmtsNode* stmts = new StmtsNode();

    for (auto it = n->node_list()->begin(); 
        it != n->node_list()->end(); it++) {
        Node* new_node = visit(*it);
        if (new_node) {
            new_node->set_silent((*it)->silent());
            stmts->add(new_node);
        }
    }

    return stmts;
}

Node* TypedVisitor::visit(ConstInt* n) {
    n->set_type(IntType::get_instance());
    return n;
}

Node* TypedVisitor::visit(ConstDouble* n) {
    n->set_type(DoubleType::get_instance());
    return n;
}

Node* TypedVisitor::visit(BuiltinConst* n) {
    switch (n->_builtin_id) {
    case BuiltinConst::UNIT:
        n->set_type(UnitType::get_instance());
        return n;
    case BuiltinConst::TRUE:
    case BuiltinConst::FALSE:
        n->set_type(BoolType::get_instance());
        return n;
    default:
        cout << "unknown const: " << n->_builtin_id;
    }
    return n;
}

Node* TypedVisitor::visit(ConstString* n) {
    n->set_type(StringType::get_instance());
    return n;
}

Node* TypedVisitor::visit(ConstArray* n) {
    if (n->arr()->size() == 0) {
        n->set_type(new ListType(UnitType::get_instance()));
        return n;
    }

    
    Type* elem_type = nullptr;
    for (int i = 1; i < n->arr()->size(); i++) {
        visit(n->arr()->get(i));
        elem_type = n->arr()->get(i)->get_type(type_table());
    }
    n->set_type(new ListType(elem_type));
    return n;
}

// TODO: add import type into symtable.
Node* TypedVisitor::visit(ImportNode* n) {
    return n;
}

Node* TypedVisitor::visit(ImplementNode* n) {
    string trait_name = n->trait_name();

    Type* tt = nullptr;
    if (trait_name != "") {
        tt = type_table()->get_expanded_type(trait_name);
    }

    TraitType* trait_type = nullptr;
    if (tt && tt->is_trait()) {
        trait_type = dynamic_cast<TraitType*>(tt);
    }

    string for_name = n->for_name();
    Type* for_type = type_table()->get_expanded_type(for_name);
    if (!for_type) {
        cerr << "TypedVisitor : undeclared record '" << n->for_name() << "'." << endl;
        _has_error = true;
        return n;
    }

    if (for_type->is_record()) {
        RecordType* record_type = dynamic_cast<RecordType*>(for_type);
        type_table()->add_substitution("Self", record_type);
    }
    else if (for_type->is_type_func()) {
        TypeFunction* type_func = dynamic_cast<TypeFunction*>(for_type);
        for_type = type_func->body();
    }

    RecordDefNode* record_def_node = dynamic_cast<RecordDefNode*>(
        _sym_table->get_subtree_by_name(for_name));

    if (!for_type->is_record() || !record_def_node) {
        cerr << "TypedVisitor : can not find implementation of '" << n->for_name() << "'." << endl;
        _has_error = true;
        return n;
    }

    _sym_table = record_def_node->sym_table();
    RecordType* record_type = dynamic_cast<RecordType*>(for_type);
    type_table()->add_substitution("Self", record_type);

    for (auto it = n->methods()->node_list()->begin();
        it != n->methods()->node_list()->end(); it++) {
        FuncDefNode* func_def_node = dynamic_cast<FuncDefNode*>(visit(*it));
        add_method_to_record(record_def_node, func_def_node, trait_type);
    }
    _sym_table = _sym_table->parent();
    
    return n;
}

void TypedVisitor::add_method_to_record(RecordDefNode* record_def_node, 
    FuncDefNode* func_def_node, TraitType* trait_type) {
    if (!func_def_node) {
        cerr << "TypedVisitor : method '" << func_def_node->name() << "' is not a function." << endl;
        _has_error = true;
        return;
    }

    if (trait_type && !trait_type->has_method(func_def_node->name())) {
        cerr << "TypedVisitor : method '" << func_def_node->name() << "' is not in trait '" << trait_type->name() << "'." << endl;
        _has_error = true;
        return;
    }

    Type* rt = record_def_node->get_type(type_table());
    RecordType* record_type = nullptr;
    // record normal_type {}
    if (rt->is_record()) {
        record_type = dynamic_cast<RecordType*>(rt);
    }
    // record generic_type[T] {}
    else if (rt->is_type_func()) {
        TypeFunction* type_func = dynamic_cast<TypeFunction*>(rt);
        record_type = dynamic_cast<RecordType*>(type_func->body());
    }
    assert(record_type);

    _sym_table = func_def_node->sym_table();
    Type* func_type = build_func_type(func_def_node);

    // To skip first parameter —— self:Self
    assert(func_type && func_type->is_arrow());
    ArrowType* arrow_type = dynamic_cast<ArrowType*>(func_type);
    func_type = arrow_type->try_apply(record_type);
    assert(func_type);
    if (!func_type->is_arrow()) {
        func_type = new ArrowType(UnitType::get_instance(), func_type);
    }

    // add symbol to Record SymbolTable
    func_def_node->set_type(func_type);
    _sym_table = _sym_table->parent();
    _sym_table->set_local_type(func_def_node->name(), func_type);
    
    record_def_node->add_method(func_def_node);
    record_type->add_method_type(func_def_node->name(), func_type);
}

Node* TypedVisitor::visit(IfNode* n) {
    IfNode* new_node = new IfNode();
    new_node->set_cond(visit(n->cond()));
    new_node->set_then_part(visit(n->then_part()));

    if (n->else_part() != NULL) {
        new_node->set_else_part(visit(n->else_part()));
    }
    return new_node;
}

Node* TypedVisitor::visit(WhileNode* n) {
    WhileNode* new_node = new WhileNode();
    new_node->set_cond(visit(n->cond()));
    new_node->set_body(visit(n->body()));
    return new_node;
}

Node* TypedVisitor::visit(CmpNode* n) {
    CmpNode* new_node = new CmpNode(n->cmp_op());
    new_node->set_left(visit(n->left()));
    new_node->set_right(visit(n->right()));
    new_node->set_type(BoolType::get_instance());
    return new_node;
}

Node* TypedVisitor::visit(LogicNotNode* n) {
    LogicNotNode* new_node = new LogicNotNode(visit(n->value()));
    new_node->set_type(BoolType::get_instance());
    return new_node;
}

Node* TypedVisitor::visit(AssignNode* n) {
    AssignNode* new_node = new AssignNode(
        visit(n->left()),
        visit(n->right())
    );

    return new_node;
}

Node* TypedVisitor::visit(VarNode* n) {
    int index = _sym_table->get_name_index(n->name());
    if (index >= 0) {
        n->set_type(_sym_table->get_name_type(n->name()));
        return n;
    }

    index = _sym_table->get_var_name_index(n->name());
    if (index >= 0) {
        n->set_type(_sym_table->get_var_name_type(n->name()));
        return n;
    }

    index = _sym_table->get_local_index(n->name());
    if (index >= 0) {
        n->set_type(_sym_table->get_local_type(n->name()));
        return n;
    }

    index = _sym_table->get_cell_var_index(n->name());
    if (index >= 0) {
        n->set_type(_sym_table->get_cell_var_type(n->name()));
        return n;
    }

    index = _sym_table->get_free_var_index(n->name());
    if (index >= 0) {
        n->set_type(_sym_table->get_free_var_type(n->name()));
        return n;
    }

    // find name along the sym table chain.
    SymbolTable* pst = _sym_table->parent();
    while (pst != nullptr) {
        index = pst->get_name_index(n->name());
        if (index >= 0) {
            n->set_type(pst->get_name_type(n->name()));
            break;
        }

        if (!pst->is_func()) {
            pst = pst->parent();
            continue;
        }

        // If pst is func, try to find variable in varnames.
        index = pst->get_var_name_index(n->name());
        if (index >= 0) {
            n->set_type(pst->get_var_name_type(n->name()));
            add_free_var(pst, n->name(), n->get_type(type_table()));
            break;
        }

        index = pst->get_local_index(n->name());
        if (index >= 0) {
            n->set_type(pst->get_local_type(n->name()));
            add_free_var(pst, n->name(), n->get_type(type_table()));
            break;
        }
        pst = pst->parent();
    }

    if (index == -1) {
        Type* vt = type_table()->get_expanded_type(n->name());
        if (!vt) {
            cerr << "TypedVisitor :undeclared variable '" << n->name() << "'." << endl;
            _has_error = true;
        }
        n->set_type(vt);
    }

    return n;
}

void TypedVisitor::add_free_var(SymbolTable* pst, const string& name, Type* type) {
    /* For those functions that variable is not defined in,
     * this variable is free, else it is cell.
     */
    SymbolTable* begin = _sym_table;
    while (begin != pst) {
        if (begin->get_free_var_index(name) < 0) {
            begin->add_free_var(name, type);
        }
        begin = begin->parent();
    }

    if (pst->get_cell_var_index(name) < 0) {
        pst->add_cell_var(name);
    }
}

Node* TypedVisitor::visit(DotNode* n) {
    return n;
}

Node* TypedVisitor::visit(SequenceNode* n) {
    visit(n->left());
    visit(n->right());
    return n;
}

/*
 * All local variables are put in varnames.
 */
void TypedVisitor::check_name(const string& name, Type* t) {
    if (_sym_table->is_func()) {
        int index = -1;
        index = _sym_table->get_var_name_index(name);
        if (index != -1) {
            cerr << "TypedVisitor : redefination for param variable '" << name << "'." << endl;
            _has_error = true;
            return;
        }

        index = _sym_table->get_local_index(name);
        if (index != -1) {
            cerr << "TypedVisitor : redefination for local variable '" << name << "'." << endl;
            _has_error = true;
            return;
        }

        _sym_table->add_local(name, t);
    }
    else {
        int index = _sym_table->get_name_index(name);
        if (index != -1) {
            cerr << "TypedVisitor : redefinition for variable '" << name << "'." << endl;
            _has_error = true;
            return;
        }
        _sym_table->add_name(name, t);
    }
}

void TypedVisitor::update_name_type(const string& name, Type* t) {
    if (_sym_table->is_func()) {
        int index = -1;
        index = _sym_table->get_var_name_index(name);
        if (index != -1) {
            _sym_table->set_var_name_type(name, t);
            return;
        }

        index = _sym_table->get_local_index(name);
        if (index != -1) {
            _sym_table->set_local_type(name, t);
            return;
        }
    }

    int index = _sym_table->get_name_index(name);
    if (index != -1) {
        _sym_table->set_name_type(name, t);
        return;
    }

    return;
}

Node* TypedVisitor::visit(CallNode* n) {
    Node* callee = visit(n->callee());
    // callee may be different from n.
    n->set_callee(callee);
    ListNode* args = dynamic_cast<ListNode*>(visit((Node*)n->args()));

    Type* callee_type = callee->get_type(type_table());
    if (!callee_type) {
        // we can not determine the type of the callee, return a type variable
        // to represent the type of the function call result.
        return n;
    }

    if (callee_type->is_list()) {
        ListType* lt = dynamic_cast<ListType*>(callee_type);
        n->set_type(lt);
        return n;
    }

    if (callee_type->is_record()) {
        RecordType* rt = dynamic_cast<RecordType*>(callee_type);
        n->set_type(rt);
        return n;
    }

    return n;
}

Node* TypedVisitor::visit(FuncDefNode* n) {
    string name = n->name();
    check_name(name, AnyType::get_instance());

    _sym_table = new SymbolTable(_sym_table, true);
    n->set_sym_table(_sym_table);

    _sym_table->set_arg_count(n->args()->size());

    if (n->has_generic_args()) {
        for (auto it = n->generic_args()->begin(); 
            it != n->generic_args()->end(); it++) {
            type_table()->add_generic_name(*it);
        }
    }

    // Return type got from declaration.
    // It should be same as declared type of function.
    if (n->ret_type_node()) {
        _sym_table->set_ret_type(n->ret_type_node()->get_type(type_table()));
    }
    else {
        _sym_table->set_ret_type(TypeVar::make_type_var(type_table()));
    }
    n->set_ret_type(_sym_table->ret_type());

    n->set_args(dynamic_cast<ListNode*>(visit((Node*)n->args())));
    // function definition may not have body.
    if (n->body()) {
        n->set_body(dynamic_cast<StmtsNode*>(visit((Node*)n->body())));
    }

    Type* func_type = build_func_type(n);
    if (n->has_generic_args()) {
        func_type = new PolyType(func_type, type_table()->get_generic_types());
    }

    _sym_table = _sym_table->parent();

    update_name_type(name, func_type);

    n->set_type(func_type);
    return n;
}

Node* TypedVisitor::visit(RecordDefNode* n) {
    string name = n->name();

    TypeFunction* type_func = nullptr;
    Type* record_type = nullptr;

    // AnyType is a placeholder for record type.
    // '''record List[T] {head : T, tail : List[T]}'''
    // in case of recursion.
    type_table()->add(name, AnyType::get_instance());
    _sym_table->set_subtree(name, n);

    _sym_table = new SymbolTable(_sym_table, true);
    n->set_sym_table(_sym_table);

    if (n->has_generic_args()) {
        type_func = new TypeFunction(nullptr);
        for (auto it = n->generic_args()->begin();
            it != n->generic_args()->end(); it++) {
            visit(*it);
            type_func->args().push_back((*it)->get_type(type_table()));
        }
    }

    map<string, Type*> fields;
    vector<Type*> field_types;

    ListNode* fields_list = new ListNode();
    for (auto it = n->fields()->node_list()->begin();
        it != n->fields()->node_list()->end(); it++) {
        VarDefNode* field_node = dynamic_cast<VarDefNode*>(visit(*it));
        fields_list->add(field_node);
        fields[field_node->name()] = field_node->get_type(type_table());
        field_types.push_back(field_node->get_type(type_table()));
    }
    n->set_fields(fields_list);

    if (name == "List") {
        record_type = new ListType(type_func->args()[0]);
    }
    else {
        record_type = new RecordType(name, fields, field_types);
    }

    if (type_func) {
        type_func->set_body(record_type);
        n->set_type(type_func);
    }
    else {
        n->set_type(record_type);
    }

    _sym_table = _sym_table->parent();
    // refresh the type of the record.
    type_table()->add(name, n->get_type(type_table()));

    return n;
}

Node* TypedVisitor::visit(TraitDefNode* n) {
    // AnyType is a placeholder for trait type.
    // '''trait List[T] {head : T, tail : List[T]}'''
    // in case of recursion.
    type_table()->add(n->name(), AnyType::get_instance());

    _sym_table = new SymbolTable(_sym_table, true);
    n->set_sym_table(_sym_table);

    type_table()->add_substitution("Self", AnyType::get_instance());

    map<string, Type*> methods;
    for (auto it = n->methods()->node_list()->begin();
        it != n->methods()->node_list()->end(); it++) {
        FuncDefNode* fd = dynamic_cast<FuncDefNode*>(visit(*it));
        if (!fd) {
            cerr << "TypedVisitor : method '" << *it << "' is not a function." << endl;
            _has_error = true;
            _sym_table = _sym_table->parent();
            return n;
        }
        methods[fd->name()] = fd->get_type(type_table());
    }
    TraitType* tt = new TraitType(n->name(), methods);
    n->set_type(tt);
    type_table()->add(n->name(), tt);
    type_table()->add_substitution("Self", tt);

    _sym_table = _sym_table->parent();
    // refresh the type of the trait.
    type_table()->add(n->name(), tt);

    return n;
}

Node* TypedVisitor::visit(SubscrNode* n) {
    visit(n->name());
    visit(n->index());

    Type* name_type = n->name()->get_type(type_table());

    if (!name_type) {
        return n;
    }
    
    if (name_type->is_poly()) {
        // Handle generic instantiation
        // The parser has already parsed this as type arguments
        ListNode* type_args = dynamic_cast<ListNode*>(n->index());
        if (!type_args) {
            cerr << "TypedVisitor : generic instantiation expects type arguments list" << endl;
            _has_error = true;
            return nullptr;
        }

        VarNode* name_node = dynamic_cast<VarNode*>(n->name());
        if (!name_node) {
            cerr << "TypedVisitor : generic function name must be a simple identifier" << endl;
            _has_error = true;
            return nullptr;
        }

        PolyType* poly_type = dynamic_cast<PolyType*>(name_type);
        if (type_args->size() != poly_type->type_args.size()) {
            cerr << "TypedVisitor : generic function type arguments mismatch" << endl;
            _has_error = true;
            return nullptr;
        }
        
        TypeTable* env = nullptr;
        for (int i = 0; i < type_args->size(); i++) {
            // must be type var.
            TypeVar* tv = dynamic_cast<TypeVar*>(poly_type->type_args[i]);
            env = tv->env();
            env->add_substitution(tv->name(), type_args->get(i)->get_type(type_table()));
        }

        Type* expanded_type = poly_type->expand(env);
        n->set_type(expanded_type);

        return n;
    }

    if (name_type->is_list()) {
        ListType* lt = dynamic_cast<ListType*>(name_type);
        Type* index_type = n->index()->get_type(type_table());
        if (index_type && index_type == IntType::get_instance()) {
            n->set_type(lt->elm_type());
        }
        return n;
    }

    // name_type has not been decided yet.
    // unify it in typechecker later.
    if (name_type->is_type_var()) {
        n->set_type(name_type);
        return n;
    }

    return n;
}

Type* TypedVisitor::get_param_type(ParamDefNode* n) {
    if (_sym_table->get_var_name_index(n->name()) != -1) {
        return _sym_table->get_var_name_type(n->name());
    }

    Type* t = nullptr;
    if (n->def_type()) {
        t = n->def_type()->get_type(type_table());
    }
    else {
        t = TypeVar::make_type_var(type_table());
    }
    return t;
}

Type* TypedVisitor::build_func_type(FuncDefNode* n) {
    ListNode* args = n->args();

    Type* func_type = n->ret_type();
    assert(func_type);

    if (args->size() == 0) {
        return new ArrowType(UnitType::get_instance(), func_type);
    }

    for (auto it = args->node_list()->rbegin();
        it != args->node_list()->rend(); it++) {
        Type* t = (*it)->get_type(type_table());
        func_type = new ArrowType(t, func_type);
    }

    return func_type;
}

Node* TypedVisitor::visit(ParamDefNode* n) {
    Type* t = get_param_type(n);
    if (_sym_table->get_var_name_index(n->name()) != -1) {
        cerr << "TypedVisitor : redefination for variable '" << n->name() << "'." << endl;
        _has_error = true;
        return nullptr;
    }

    _sym_table->add_var_name(n->name(), t);
    n->set_type(t);

    return n;
}

Node* TypedVisitor::visit(VarDefNode* n) {
    string name = n->name();
    VarDefNode* new_node = new VarDefNode(name);
    Type* declare_type = nullptr;
    Type* value_type = nullptr;

    if (n->def_type()) {
        Type* t = n->def_type()->get_type(type_table());
        if (t) {
            declare_type = t;
        }
        else {
            cerr << "undeclared type."<< endl;
            _has_error = true;
        }
    }
    else {
        declare_type = TypeVar::make_type_var(type_table());
    }

    new_node->set_type(declare_type);
    new_node->set_def_type(nullptr);

    check_name(name, declare_type);

    if (n->init_value()) {
        new_node->set_init_value(visit(n->init_value()));
    }

    return new_node;
}

Node* TypedVisitor::visit(TypeVarNode* n) {
    Type* node_type = type_table()->get_expanded_type(n->name());
    if (node_type == nullptr) {
        node_type = new TypeVar(n->name(), type_table());
        type_table()->add(n->name(), node_type);
    }

    n->set_type(node_type);

    return nullptr;
}

Node* TypedVisitor::visit(TypeVarDefNode* n) {
    visit(n->value());
    type_table()->add(n->name(), n->value()->get_type(type_table()));

    return nullptr;
}

Node* TypedVisitor::visit(TypeFuncDefNode* n) {
    TypeFunction* tf = new TypeFunction(nullptr);

    enter_type_scope();
    for (auto it = n->args()->node_list()->begin(); 
        it != n->args()->node_list()->end(); it++) {
        TypeVarNode* tvn = dynamic_cast<TypeVarNode*>(*it);
        visit(tvn);
        tf->args().push_back(tvn->get_type(type_table()));
    }

    if (n->body()) {
        visit(n->body());
        tf->set_body(n->body()->get_type(type_table()));
    }

    exit_type_scope();

    n->set_type(tf);
    type_table()->add(n->name(), tf);

    return nullptr;
}

void TypedVisitor::enter_type_scope() {
    _sym_table = new SymbolTable(_sym_table, true);
}

void TypedVisitor::exit_type_scope() {
    _sym_table = _sym_table->parent();
}

Node* TypedVisitor::visit(TypeCallNode* n) {
    return n;
}

Node* TypedVisitor::visit(TypeConstInt* n) {
    n->set_type(new TypeDependentConst(n->value));

    return nullptr;
}

Node* TypedVisitor::visit(ArrowTypeNode* n) {
    visit((Node*)n->src());
    Type* left = n->src()->get_type(type_table());
    visit((Node*)n->dst());
    Type* right = n->dst()->get_type(type_table());
    n->set_type(new ArrowType(left, right));

    return nullptr;
}

Node* TypedVisitor::visit(LambdaDef* n) {
    _sym_table = new SymbolTable(_sym_table, true);
    n->set_sym_table(_sym_table);

    _sym_table->set_arg_count(1);
    _sym_table->set_ret_type(TypeVar::make_type_var(type_table()));

    visit(n->param());
    Type* param_type = n->param()->get_type(type_table());
    visit(n->body());
    n->set_type(new ArrowType(param_type, _sym_table->ret_type()));

    /* Restore to parent's scope */
    _sym_table = _sym_table->parent();

    return n;
}

Node* TypedVisitor::visit(TupleNode* n) {
    ListNode* args = n->args();
    vector<Type*> types;
    for (auto it = args->node_list()->begin();
        it != args->node_list()->end(); it++) {
        Type* t = (*it)->get_type(type_table());
        types.push_back(t);
    }
    n->set_type(new TupleType(types));
    return n;
}

Node* TypedVisitor::visit(TypeTupleNode* n) {
    return nullptr;
}

/* pseudo nodes */
Node* TypedVisitor::visit(ArrayTypeNode* n) {
    return nullptr;
}
